#%%
### This script was written by Eric Deal, 06.22.2020
### Script should be run with python 3
%matplotlib inline
import numpy as np, pandas as pd, pathlib, matplotlib.pyplot as plt

df = {}
for fn in list(pathlib.Path('observations/4_data_settling_velocities/').glob('*_settling_velocities.csv')):
    df_temp = pd.read_csv(fn)
    df[fn.stem.split('_')[0]] = {'vel_mean': df_temp.vel.mean()/1e3, 'dvel_mean': df_temp.vel.std()/1e3/np.sqrt(df_temp.vel.size), 'vels': df_temp.vel.values/1e3, 'dvels': df_temp.vel_e.values/1e3}


df = pd.DataFrame(df).T
#%%
df_vol = pd.read_csv('1_var_density.csv').set_index('Unnamed: 0').rename_axis(None)
df_shape = pd.read_csv('2_var_shape.csv').set_index('Unnamed: 0').rename_axis(None)

#%%
df['CD'] = np.nan
for grain in df.vel_mean.index:
    df['CD'][grain] = 4*(df_vol.total_grain_ro[grain]/990 - 1)*9.81*df_shape.dn[grain] / (3*df.vel_mean[grain]**2) * df_shape.CSF[grain]

df.to_csv('4_var_settling_velocities.csv')
#%%
df = df.T
fig, axs = plt.subplots(1,5, figsize=(10, 3), sharey=True)
label = {'gb': 'Spheres', 'oc': 'Faceted ellipsoids', 'gg': 'Rounded chips', 'ls': 'Rectangular prisms', 'ng': 'Natural gravel'}
ns = np.zeros(5)
for i, grain in enumerate(['gb', 'oc', 'gg', 'ng', 'ls']):
        axs[i].hist(df[grain]['vels'], density=True, alpha=.6, bins=20, range=(0,.5), color='k') # A
        ns[i] = df[grain]['vels'].size
        axs[i].set_title(label[grain] + '\n(n = %i)' % ns[i])

for i in range(5):
        axs[i].set_xlim(0.2,0.6)
       
axs[2].set_xlabel('Settling velocity (m/s)', fontsize=14)


axs[0].set_ylabel('# observations', fontsize=14)

plt.tight_layout()
plt.savefig('Grain_settling_dists.png')

# %%
